package JzOffer.jz47;

import java.util.ArrayDeque;
import java.util.Deque;

/**
 * @author 张取
 * @version 1.0
 * @date 2022/1/20 下午10:39
 * JzOffer.jz47
 */

public class Solution02 {
    public TreeNode pruneTree(TreeNode root){
        if(root == null){
            return null;
        }
        Deque<TreeNode> stack = new ArrayDeque<>();
        TreeNode curNode = root;
        TreeNode prevNode = null;
        while(curNode != null || !stack.isEmpty()){
            while(curNode != null){
                stack.offerLast(curNode);
                curNode = curNode.left;
            }
            curNode = stack.peekLast();
            if(curNode.right != null && curNode.right != prevNode) {
                curNode = curNode.right;
           } else {
                curNode = stack.pollLast();
                if(curNode.val == 0 && curNode.left == null && curNode.right == null){
                    //curNode为栈区变量，直接修改curNode对堆区的树没有影响
                    //curNode = null;

                    //tmpNode为栈区变量,tmpNode.left和tmpNode.right为堆区变量，修改left和right会修改树的节点
                    if(!stack.isEmpty()){
                        //tmpNode为父节点,curNode为子节点
                        TreeNode tmpNode = stack.peekLast();
                        if(tmpNode.left == curNode){
                            tmpNode.left = null;
                        } else if(tmpNode.right == curNode){
                            tmpNode.right = null;
                        }
                    } else {
                        //出栈前栈中只有一个节点，直接返回
                        return null;
                    }
                }
                prevNode = curNode;
                curNode = null;
            }
        }
        return root;
    }
}
